import time

import cv2
import numpy as np
from PIL import Image, ImageEnhance

MORPH_MAP = {
    'MORPH_BLACKHAT': cv2.MORPH_BLACKHAT,
    'MORPH_CLOSE': cv2.MORPH_CLOSE,
    'MORPH_CROSS': cv2.MORPH_CROSS,
    'MORPH_DILATE': cv2.MORPH_DILATE,
    'MORPH_ELLIPSE': cv2.MORPH_ELLIPSE,
    'MORPH_ERODE': cv2.MORPH_ERODE,
    'MORPH_GRADIENT': cv2.MORPH_GRADIENT,
    'MORPH_HITMISS': cv2.MORPH_HITMISS,
    'MORPH_OPEN': cv2.MORPH_OPEN,
    'MORPH_RECT': cv2.MORPH_RECT,
    'MORPH_TOPHAT': cv2.MORPH_TOPHAT
}


MORPH_SHAPE = {
    'MORPH_CROSS': cv2.MORPH_CROSS,
    'MORPH_ELLIPSE': cv2.MORPH_ELLIPSE,
    'MORPH_RECT': cv2.MORPH_RECT,
    'MORPH_HLINE': None,
    'MORPH_VLINE': None
}

THRESHOLD_TYPES = {
    'THRESH_BINARY': cv2.THRESH_BINARY,
    'THRESH_BINARY_INV': cv2.THRESH_BINARY_INV
}

ADAPTIVE_THRESH_TYPES = {
    'ADAPTIVE_THRESH_MEAN_C': cv2.ADAPTIVE_THRESH_MEAN_C,
    'ADAPTIVE_THRESH_GAUSSIAN_C': cv2.ADAPTIVE_THRESH_GAUSSIAN_C
}


# 图像预处理
def gaussian_blur(image, kernel_size):
    if isinstance(kernel_size, int):
        kernel_size = (kernel_size, kernel_size)
    gauss_blur = cv2.GaussianBlur(image, kernel_size, 0)
    return gauss_blur


def median_blur(image, kernel_size):
    result = cv2.medianBlur(image, kernel_size)
    return result


def adjust_color(image, factor):
    if len(image.shape) == 3:
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    pil_image = Image.fromarray(image, mode='RGB')
    enhancer = ImageEnhance.Color(pil_image)
    pil_image = enhancer.enhance(factor)
    return np.array(pil_image, dtype=image.dtype)[..., ::-1]


def adjust_contrast(image, factor):
    assert image.dtype == np.uint8, 'Pillow backend only support uint8 type'
    # Image.fromarray defaultly supports RGB, not BGR.
    if len(image.shape) == 3:
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    pil_image = Image.fromarray(image, mode='RGB')
    enhancer = ImageEnhance.Contrast(pil_image)
    pil_image = enhancer.enhance(factor)
    return np.array(pil_image, dtype=image.dtype)[..., ::-1]


def adjust_brightness(image, factor):
    assert image.dtype == np.uint8, 'Pillow backend only support uint8 type'
    # Image.fromarray defaultly supports RGB, not BGR.
    if len(image.shape) == 3:
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    pil_image = Image.fromarray(image, mode='RGB')
    enhancer = ImageEnhance.Brightness(pil_image)
    pil_image = enhancer.enhance(factor)
    return np.array(pil_image, dtype=image.dtype)[..., ::-1]


def adjust_sharpness(image, factor):
    assert image.dtype == np.uint8, 'Pillow backend only support uint8 type'
    # Image.fromarray defaultly supports RGB, not BGR.
    if len(image.shape) == 3:
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    pil_image = Image.fromarray(image, mode='RGB')
    enhancer = ImageEnhance.Sharpness(pil_image)
    pil_image = enhancer.enhance(factor)
    return np.array(pil_image, dtype=image.dtype)[..., ::-1]


# def adjust_sharpness(img, factor=1., kernel=None):
#     """Adjust image sharpness."""
#     if kernel is None:
#         # adopted from PIL.ImageFilter.SMOOTH
#         kernel = np.array([[1., 1., 1.], [1., 5., 1.], [1., 1., 1.]]) / 13
#     assert isinstance(kernel, np.ndarray), \
#         f'kernel must be of type np.ndarray, but got {type(kernel)} instead.'
#     assert kernel.ndim == 2, \
#         f'kernel must have a dimension of 2, but got {kernel.ndim} instead.'
#
#     degenerated = cv2.filter2D(img, -1, kernel)
#     sharpened_img = cv2.addWeighted(
#         img.astype(np.float32), factor, degenerated.astype(np.float32),
#         1 - factor, 0)
#     sharpened_img = np.clip(sharpened_img, 0, 255)
#     return sharpened_img.astype(img.dtype)


def gamma_transform(img, gamma):
    if img.dtype == np.uint8:
        table = (np.arange(0, 256.0 / 255, 1.0 / 255) ** gamma) * 255
        img = cv2.LUT(img, table.astype(np.uint8))
    else:
        img = np.power(img, gamma)

    return img


def equalize(image):
    if len(image.shape) == 2 or image.shape[2] == 1:
        image = cv2.equalizeHist(image)
    else:
        image = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
        image[:, :, 0] = cv2.equalizeHist(image[:, :, 0])
        image = cv2.cvtColor(image, cv2.COLOR_LAB2BGR)

    return image


def clahe(img, clip_limit=2.0, tile_grid_size=8):
    if img.dtype != np.uint8:
        raise TypeError("clahe supports only uint8 inputs")

    if isinstance(tile_grid_size, int):
        tile_grid_size = [tile_grid_size] * 2

    clahe_mat = cv2.createCLAHE(
        clipLimit=clip_limit, tileGridSize=tile_grid_size)

    if len(img.shape) == 2 or img.shape[2] == 1:
        img = clahe_mat.apply(img)
    else:
        img = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
        img[:, :, 0] = clahe_mat.apply(img[:, :, 0])
        img = cv2.cvtColor(img, cv2.COLOR_LAB2BGR)

    return img


def detail_enhance(image, sigma_s, sigma_r):
    result = cv2.detailEnhance(image, None, sigma_s, sigma_r)
    return result


# 形态学操作
def morphology(image, kernel_size, kernel_shape, method):
    if len(image.shape) == 3:
        image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    if isinstance(kernel_size, int):
        kernel_size = [kernel_size] * 2
    if kernel_shape == "MORPH_HLINE":
        kernel = np.ones((1, kernel_size[0]), dtype=np.uint8)
    elif kernel_shape == "MORPH_VLINE":
        kernel = np.ones((kernel_size[0], 1), dtype=np.uint8)
    else:
        kernel_shape = MORPH_MAP[kernel_shape]
        kernel = cv2.getStructuringElement(kernel_shape, kernel_size)
    method = MORPH_MAP[method]
    result = cv2.morphologyEx(image, method, kernel)
    return result


def morphology_open(image, kernel_size, kernel_shape):
    if isinstance(kernel_size, int):
        kernel_size = [kernel_size] * 2
    if kernel_shape == "MORPH_HLINE":
        kernel = np.ones((1, kernel_size[0]), dtype=np.uint8)
    elif kernel_shape == "MORPH_VLINE":
        kernel = np.ones((kernel_size[0], 1), dtype=np.uint8)
    else:
        kernel_shape = MORPH_MAP[kernel_shape]
        kernel = cv2.getStructuringElement(kernel_shape, kernel_size)
    result = cv2.morphologyEx(image, cv2.MORPH_OPEN, kernel)
    return result


def morphology_close(image, kernel_size, kernel_shape):
    if isinstance(kernel_size, int):
        kernel_size = [kernel_size] * 2
    if kernel_shape == "MORPH_HLINE":
        kernel = np.ones((1, kernel_size[0]), dtype=np.uint8)
    elif kernel_shape == "MORPH_VLINE":
        kernel = np.ones((kernel_size[0], 1), dtype=np.uint8)
    else:
        kernel_shape = MORPH_MAP[kernel_shape]
        kernel = cv2.getStructuringElement(kernel_shape, kernel_size)
    result = cv2.morphologyEx(image, cv2.MORPH_CLOSE, kernel)
    return result


# 边缘检测
def edge_sobel(image, dx, dy, ksize):
    ksize = ksize + 1 if ksize % 2 == 0 else ksize
    edge_64f = cv2.Sobel(image, cv2.CV_64F, dx, dy, ksize=ksize)
    edge_abs = np.absolute(edge_64f)
    edge_8u = np.uint8(edge_abs)
    return edge_8u


def edge_laplacian(image, ksize):
    ksize = ksize + 1 if ksize % 2 == 0 else ksize
    edge_64f = cv2.Laplacian(image, cv2.CV_64F, ksize=ksize)
    edge_abs = np.absolute(edge_64f)
    edge_8u = np.uint8(edge_abs)
    return edge_8u


def edge_canney(image, threshold1, threshold2, aperture_size):
    aperture_size = (aperture_size + 1 if aperture_size % 2 == 0
                     else aperture_size)
    edge_8u = cv2.Canny(
        image, threshold1, threshold2, apertureSize=aperture_size)
    return edge_8u


# 阈值分割
def custom_threshold(image, thresh_value, thresh_type):
    threshold_type = THRESHOLD_TYPES[thresh_type]
    print(f"threshold_type: {threshold_type}")
    if len(image.shape) == 3:
        image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    binary = cv2.threshold(image, thresh_value, 255, threshold_type)[-1]
    return binary


def adaptive_threshold(image, block_size, constant, thresh_type,
                       adaptive_thresh_type):
    if len(image.shape) == 3:
        image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    threshold_type = THRESHOLD_TYPES[thresh_type]
    adaptive_thresh_type = ADAPTIVE_THRESH_TYPES[adaptive_thresh_type]

    binary = cv2.adaptiveThreshold(image,
                                   255,
                                   adaptive_thresh_type,
                                   threshold_type,
                                   block_size,
                                   constant)
    return binary


if __name__ == '__main__':
    file = r"F:\coco\train\2023_04_10_16_15_32_885_4.bmp"
    image = cv2.imdecode(np.fromfile(file, dtype=np.uint8), 1)
    # edge_sobel(image, 1, 1, 3)
    # edge_laplacian(image, 3)
    edge_canney(image, 3, 8, 5)




